package com.jsean.web;

import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.java_websocket.client.WebSocketClient;
import org.java_websocket.exceptions.WebsocketNotConnectedException;
import org.java_websocket.handshake.ServerHandshake;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.net.URI;
import java.util.Base64;

@Slf4j
public class FileTransferClient extends WebSocketClient {
    public FileTransferClient(URI serverUri) {
        super(serverUri);
    }


    private static final int CHUNK_SIZE = 1024 * 64; // 64KB chunks
    private String fileName;
    private long fileSize;
    private int sequence = 0;
    private String KEYWORD = "^^";
    private String SPLIT_WORD = "\\^\\^";

    public void uploadFile(String filePath) {
        try {
            while (!this.isOpen()) {
                log.info("Connection not established yet");
                Thread.sleep(1000L);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }

        try {
            File file = new File(filePath);
            this.fileSize = file.length();
            this.fileName = filePath;
            this.sequence = 0;
            
            // Send file metadata first
            String metadata = "FILE_METADATA" + KEYWORD + fileName + KEYWORD + fileSize;
            try {
                this.send(metadata);
            } catch (WebsocketNotConnectedException e) {
                log.info("Failed to send - connection not ready");
            }

            // Then send file chunks
            FileInputStream fis = new FileInputStream(file);
            byte[] buffer = new byte[CHUNK_SIZE];
            int bytesRead;
            
            while ((bytesRead = fis.read(buffer)) != -1) {
                String chunk = "CHUNK" + KEYWORD + sequence + KEYWORD +
                              Base64.getEncoder().encodeToString(buffer);
                this.send(chunk);
                sequence++;
            }
            
            fis.close();
            this.send("FILE_END" + KEYWORD + fileName);
            log.info("File sent in " + sequence + " chunks");

        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    @Override
    public void onOpen(ServerHandshake handshakedata) {
        log.info("Connected to server");
    }

    private FileOutputStream downloadFileStream;
    private long receivedFileSize = 0;
    private long expectedDownloadSize = 0;
    private String downloadingFileName;

    @Override
    public void onMessage(String message) {
        try {
            if (message.startsWith("FILE_METADATA" + KEYWORD)) {
                // Handle download file metadata
                String[] parts = message.split(SPLIT_WORD);
                downloadingFileName = "downloaded_" + new File(parts[1]).getName();
                expectedDownloadSize = Long.parseLong(parts[2]);
                receivedFileSize = 0;

                downloadFileStream = new FileOutputStream(downloadingFileName);
                log.info("Preparing to download: " + downloadingFileName);

            } else if (message.startsWith("CHUNK" + KEYWORD)) {
                // Handle download chunks
                String[] parts = message.split(SPLIT_WORD);
                byte[] chunkData = Base64.getDecoder().decode(parts[2]);

                downloadFileStream.write(chunkData);
                receivedFileSize += chunkData.length;

                // Print progress every 10 chunks
                if (Integer.parseInt(parts[1]) % 10 == 0) {
                    double progress = (double)receivedFileSize / expectedDownloadSize * 100;
                    String formattedProgress = String.format("%.2f", progress);
                    log.info("Download progress: {}%", formattedProgress);
                }

            } else if (message.startsWith("FILE_END" + KEYWORD)) {
                // Finalize download
                downloadFileStream.close();
                log.info("File downloaded successfully: " + downloadingFileName);

                // Reset state
                downloadingFileName = null;
                expectedDownloadSize = 0;
                receivedFileSize = 0;
                downloadFileStream = null;

            } else {
                log.info("Received from server: " + message);
            }
        } catch (IOException e) {
            e.printStackTrace();
            try {
                if (downloadFileStream != null) {
                    downloadFileStream.close();
                }
            } catch (IOException ex) {
                ex.printStackTrace();
            }
        }
    }

    @Override
    public void onClose(int code, String reason, boolean remote) {
        log.info("Connection closed: " + reason);
    }

    @Override
    public void onError(Exception ex) {
        ex.printStackTrace();
    }

    public void requestFileDownload(String fileName) {
        try {
            while (!this.isOpen()) {
                log.info("Connection not established yet");
                Thread.sleep(1000L);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }

        this.send("DOWNLOAD_REQUEST" + KEYWORD + fileName);
    }

    public static void main(String[] args) {
        try {
            FileTransferClient client = new FileTransferClient(new URI("ws://localhost:8090"));
            client.connect();
            
            // Example usage:
            String fileName = "D:\\02.Projects\\java_projects\\JSean\\webSocketServer\\src\\main\\resources\\file.txt";
            if (args.length != 0) {
                fileName = args[1];
            }

//            client.uploadFile(fileName);
            client.requestFileDownload(fileName);
            client.close();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}